
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import sys
sys.path.append("models/")

sys.path.append("utils")
from class_model import VGG19
from ghostnetv2_torch import ghostnetv2
import time
import os
from math import cos, pi
import numpy as np
import cv2

from torch.utils.tensorboard import SummaryWriter #tensorboard 可视化
# %matplotlib inline
# **************************************************************
import argparse


from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import seaborn as sns  # 统计数据可视化
import matplotlib.pyplot as plt  # 数据可视化
import pandas as pd

def make_parser():
    parser = argparse.ArgumentParser(description="training config", add_help= False)

    parser.add_argument("--weight", default='output/take_photo/ghostnetv2_224_224_classify_3class/2024_01_22_11_27_20/best.pth', type=str) #ghostnet2 or 
    parser.add_argument("--image_root", default='/home/cheng/data/workspace/part-time/ear/classify/data', type=str)
    parser.add_argument("--num_classes", default=3, type=int)
    parser.add_argument("--max_epoch", default=200, type=int)
    parser.add_argument("--batch_size", default=20, type=int)
    parser.add_argument("--lr_min", default=0.0001, type=float)
    parser.add_argument("--lr_max", default=0.01, type=float)
    parser.add_argument("--input_size", default=[576, 576], type=int, nargs='+')

    parser.add_argument("--load_train", default='', type=str)
    parser.add_argument("--resume_train", default='', type=str)
    parser.add_argument("--mean", default=[0.47627547, 0.3684924, 0.32831427], type=float, nargs='+')
    parser.add_argument("--std", default=[0.32356596, 0.26539913, 0.2421861], type=float, nargs='+')
    parser.add_argument("--num_workers", default=4, type=int)
    parser.add_argument("--base_lr", default=0.0001, type=float)

    parser.add_argument("--output_name", default='', type=str)
    return parser


# **************************************************************
def adjust_learning_rate(optimizer, current_epoch, max_epoch, lr_min=0., lr_max=0.01, warmup=True):
    warmup_epoch = 5 if warmup else 0
    lr = optimizer.param_groups[0]['lr']
    lr_min = lr * 0.5
    if current_epoch < warmup_epoch:
        lr = lr_max * current_epoch / warmup_epoch
    else:
        lr = lr_min + 0.5 * (lr_max - lr_min) * (
            1.0
            + cos(
                pi
                * (current_epoch - warmup_epoch)
                / (max_epoch - warmup_epoch)
            )
        )
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
        print("lr:", lr)



def main(args):

    current_time = time.localtime()
    savefolder = os.path.join("output/" + args.output_name, time.strftime("%Y_%m_%d_%H_%M_%S", current_time))
    os.makedirs(savefolder, exist_ok = True)

    # 利用torchvision对图像数据预处理
    from preproc import static_resize
    train_transform = transforms.Compose([
                            static_resize(args.input_size),
                            # transforms.Resize(96),
                            # transforms.RandomResizedCrop(80),
                            transforms.RandomAffine(degrees=15,scale=(0.8,1.5)),
                            transforms.ToTensor(),
                            transforms.Normalize(args.mean, args.std)
                            ])

    val_transform = transforms.Compose([
                                        static_resize(args.input_size),
                                        # transforms.Resize(96),
                                        # transforms.CenterCrop(80),
                                        transforms.ToTensor(),
                                        transforms.Normalize(args.mean, args.std)
                                        ])

    trainset = torchvision.datasets.ImageFolder(root=args.image_root + '/train', transform=train_transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)

    valset = torchvision.datasets.ImageFolder(root=args.image_root + '/val', transform=val_transform)
    valloader = torch.utils.data.DataLoader(valset, batch_size=1, shuffle=False, num_workers=args.num_workers)

    # 展示训练样本和测试样本数
    print("len(trainloader) = ", len(trainloader))
    print(trainset.class_to_idx)
    print("len(valloader) = ", len(valloader))
    print(valset.class_to_idx)



    # CPU 或者 GPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    # 初始化网络,加载预训练模型
    # model = VGG19(num_classes=2, init_weights=True)
    model = ghostnetv2(num_classes=args.num_classes, 
                        width=1, 
                        dropout=0.2)

    model.load_state_dict(torch.load(args.weight, map_location='cuda:0')['model'])


    # 查看GPU可用情况
    # if torch.cuda.device_count()>1:
    #     print('We are using',torch.cuda.device_count(),'GPUs!')
    #     model = nn.DataParallel(model)
    model.to(device)

    # 训练
    since = time.time()

    acc = 0.
    model.eval()
    print('waitting for Val...')
    label_list = []
    prediction_list = []
    with torch.no_grad():
        accuracy = 0.
        total =0
        for data in valloader:
            images, labels = data
            images = images.to(device)
            labels = labels.to(device)
            out = model(images)
            _, prediction = torch.max(out, 1)
            # print("label = ", labels)
            # print("prediction = ", prediction)
            label_list.extend(list(np.array(labels.cpu())))
            prediction_list.extend(list(np.array(prediction.cpu())))
            total += labels.size(0)
            accuracy += (prediction == labels).sum().item()
        acc = 100.*accuracy/total
    # tblogger.add_scalar("val_accuracy", acc, epoch + 1) #tensorboard 可视化
    # print('epoch {}  The ValSet accuracy is {:.4f}% \n'.format(epoch + 1, acc))
    # Val_Accuracy.append(acc)

    # *****************************calculate results*************************************
    labels = np.array(label_list)
    prediction = np.array(prediction_list)
    print("precision :", precision_score(labels, prediction, average='micro'), "\n")
    print("Recall :", recall_score(labels, prediction, average='micro'), "\n")
    print("f1 score:", f1_score(labels, prediction, average='micro'), "\n")

# ROC CURVE
    from sklearn.metrics import roc_curve, roc_auc_score
    fpr, tpr, thresholds = roc_curve(labels, prediction)
    auc_score = roc_auc_score(labels, prediction)
    plt.plot(fpr, tpr, label=f'AUC = {auc_score:.2f}')  # 绘制ROC曲线，标注AUC的值
    # 随即分类器没有分类能力，其FPR=TPR。随机分类器的性能通常表示为ROC曲线上的对角线
    plt.plot([0, 1], [0, 1], linestyle='--', color='r', label='Random Classifier')  # 绘制随机分类器的ROC曲线
    plt.xlabel('False Positive Rate')  # x轴标签为FPR
    plt.ylabel('True Positive Rate')   # y轴标签为TPR
    plt.title('ROC Curve')             # 设置标题
    plt.legend()                    
    plt.show()
    # 混淆矩阵
    cm = confusion_matrix(labels, prediction, labels=[0, 1, 2])
    print('Confusion matrix\n\n', cm)
    print('\nTrue Positives(TP) = ', cm[0, 0])
    print('\nTrue Negatives(TN) = ', cm[1, 1])
    print('\nFalse Positives(FP) = ', cm[0, 1])
    print('\nFalse Negatives(FN) = ', cm[1, 0])

    # cm_matrix = pd.DataFrame(data=cm, columns=['Actual Positive:1', 'Actual Negative:0'],
    #                         index=['Predict Positive:1', 'Predict Negative:0'])

    # sns.heatmap(cm_matrix, annot=True, fmt='d', cmap='YlGnBu')

    
    # plt.show()

    from sklearn.metrics import classification_report

    # 分类报告
    print('******************************************')
    print(classification_report(labels, prediction))
    print('******************************************')
    

    # if acc > BEST_VAL_ACC:
    #     print('Find Better Model and Saving it...')
    #     ckpt_state = {
    #         "start_epoch":epoch + 1,
    #         "model": model.state_dict(),
    #         "optimizer": optimizer.state_dict(),
    #     }
    #     torch.save(ckpt_state, savefolder + "/best.pth")
    #     BEST_VAL_ACC = acc
    #     print('Saved!')
    
    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed%60))
    print('val Acc is {:.4f}%'.format(acc))

if __name__ == "__main__":
    args = make_parser().parse_args()
    main(args)
